import pygame

class QuadTree(object): #Good for now, but might get a bit complicated and/or laggy after lots of objects move through.
    #root = ''
    """Fancy smancy magic stuffs"""
    def __init__(self, x,y, width, height, root = False, parent = 0):
        self.upleft = 0
        self.upright = 0
        self.downleft = 0
        self.downright = 0
       
        self.rect = pygame.Rect(x,y, width, height)
        self.entities = []
        #self.nonCollideEntities = [] #Need to implement non collide entities for practicality.
        self.parent = parent
        self.root = root
        self.drawColor = (0,0,0)
        self.capacity = 50
        #if(self.root == True):
            #QuadTree.root = self
        
    def insertEntity(self, entity):
        self.getRoot().insertEntityTrue(entity)

    def getRoot(self):
        if(self.parent == 0):
            return self
        else:
            return self.parent.getRoot()

    def insertEntityTrue(self, entity):
        if self.root == True and not self.rect.contains(entity.rect):
            self.expand(entity)
            self.parent.insertEntityTrue(entity)

        else:

            if self.rect.colliderect(entity.rect):
                 if(self.upleft == 0):
                    self.entities.append(entity)
                    if(len(self.entities) > self.capacity):
                        self.createSubQuads()
                        tempEnts = self.entities
                        self.entities = 0
                        for entity in tempEnts:
                            self.insertEntityTrue(entity)
                 else:
                     self.upleft.insertEntityTrue(entity)
                     self.upright.insertEntityTrue(entity)
                     self.downleft.insertEntityTrue(entity)
                     self.downright.insertEntityTrue(entity)

    def expand(self, entity):
        xOffset = 0
        yOffset = 0
        if entity.rect.topleft[0] < self.rect.topleft[0]: ## This doesn't work
            xOffset = -1
        if entity.rect.topright[0] > self.rect.topright[0]:
            xOffset = 0

        if entity.rect.topleft[1] < self.rect.topleft[1]:
            yOffset = -1
        if entity.rect.bottomright[1] > self.rect.bottomright[1]:
            yOffset = 0

        self.parent = QuadTree(self.rect.x + self.rect.width * xOffset, self.rect.y + self.rect.height * yOffset, self.rect.width * 2, self.rect.height * 2)
        self.root = False
        self.parent.createSubQuads()
        self.parent.root = True
        #QuadTree.root = self.parent

        if xOffset == -1:
            if yOffset == 0:
                self.parent.upright = self
            if yOffset == -1:
                self.parent.bottomright = self

        if xOffset == 0:
            if yOffset == 0:
                self.parent.upleft = self
            if yOffset == -1:
                self.parent.downleft = self    

    def removeEntity(self, entity):
        self.getRoot().removeEntityTrue(entity) 

    def removeEntityTrue(self,entity): 
         if self.rect.colliderect(entity):
             if(self.upleft == 0):
                self.entities.remove(entity)
             else:   
                self.upleft.removeEntityTrue(entity)
                self.upright.removeEntityTrue(entity)
                self.downleft.removeEntityTrue(entity)
                self.downright.removeEntityTrue(entity)

         return self


    def getEntitiesNearby(self, rect):
        return self.getRoot().getEntitiesNearbyTrue(rect)

    def getEntitiesNearbyTrue(self, rect):
         returnList = []
         
         if self.rect.colliderect(rect):
            
            if(self.upleft):
                returnList.extend(self.upleft.getEntitiesNearbyTrue(rect))
            if(self.upright):
                returnList.extend(self.upright.getEntitiesNearbyTrue(rect))
            if(self.downleft):
                returnList.extend(self.downleft.getEntitiesNearbyTrue(rect))
            if(self.downright):
                returnList.extend(self.downright.getEntitiesNearbyTrue(rect))
            if(self.upleft == 0):
                returnList.extend(self.entities)

         
             
             

         return list(set(returnList))


    def createSubQuads(self):
        width = self.rect.width / 2
        height = self.rect.height / 2
        self.upleft = QuadTree(self.rect.topleft[0], self.rect.topleft[1], width, height, self)
        self.upright = QuadTree(self.rect.topleft[0] + width, self.rect.topleft[1], width, height, self)
        self.downleft = QuadTree(self.rect.topleft[0], self.rect.topleft[1] + height, width, height, self)
        self.downright = QuadTree(self.rect.topleft[0] + width, self.rect.topleft[1] + height, width, height, self)

        
    def draw(self, camera, screen):
        if not self.root:
            self.parent.draw(camera, screen)
        else:
            self.drawTo(camera, screen)

    def drawTo(self, camera, screen):
        pygame.draw.rect(screen, self.drawColor, (self.rect.topleft[0] - camera.rect.topleft[0], self.rect.topleft[1] - camera.rect.topleft[1], self.rect.width, self.rect.height), 1)
        
        if(self.upleft != 0):
            self.upleft.drawTo(camera, screen)
        if(self.upright != 0):
            self.upright.drawTo(camera, screen)
        if(self.downleft != 0):
            self.downleft.drawTo(camera, screen)
        if(self.downright != 0):
            self.downright.drawTo(camera, screen)
